import backtrader as bt
import logging


class MartingaleStrategy(bt.Strategy):
    params = dict(
        min_gap=-221000,
        stoploss=50000,
        #  stoploss=13000,
        #   targetprofit=16000
        targetprofit=190000
    )

    def __init__(self):
        # finished orders
        self.open_orders = []
        self.stop_order = None
        self.limit_order = None
        self.close_order = None

        self.closed_orders = []
        self.martins = [1]*1000

        self.point = 0.001
        self.margin = 1  # cash
        self.martin_num = 0
        self.martin_num_limit = 20
        self.min_gap_price = -22100 * self.point
        self.base_balance = self.broker.get_cash()
        # self.base_martin = self.base_balance // sum(self.martins[0:self.martin_num_limit]) // self.margin * self.margin
        self.base_martin = 50
        self.bar_executed = 0


    def log(self, txt, dt=None):
        # return
        dt = dt or self.datetime.datetime(0)
        # logging.info('%s, %s' % (dt.isoformat(), txt))
        print('%s, %s, %s' % (dt.isoformat(), txt, self.broker.get_cash()))

    def log_order(self, order):
        tojoin = list()
        tojoin.append('Ref: {}'.format(order.ref))
        tojoin.append('OrdType: {}'.format(order.ordtypename()))
        tojoin.append('ExecType: {}'.format(order.getordername()))
        tojoin.append('Status: {}'.format(order.getstatusname()))
        tojoin.append('Size: {}'.format(order.size))
        tojoin.append('Price: {}'.format(order.price))
        tojoin.append('Price Limit: {}'.format(order.pricelimit))
        tojoin.append('TrailAmount: {}'.format(order.trailamount))
        self.log(', '.join(tojoin))

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # self.log_order(order)
            return

        # Attention: broker could reject order if not enougth cash
        if order.status in [order.Completed]:
            if order in self.open_orders:
                self.log(
                    'OPEN ORDER %s EXECUTED, Price: %.3f, Size: %.3f, Cost: %.3f, Comm %.3f' %
                    (
                        order.ordtypename(),
                        order.executed.price,
                        order.executed.size,
                        order.executed.value,
                        order.executed.comm
                    )
                )

                self.bar_executed = len(self)
                self.martin_num += 1

            elif order == self.close_order:
                self.martin_num = 0
                self.log(
                    'CLOSE ORDER EXECUTED, Price: %.3f, Size: %.3f, Cost: %.3f, Comm %.3f' %
                    (
                        order.executed.price,
                        order.executed.size,
                        order.executed.value,
                        order.executed.comm
                    )
                )
            elif order in [self.stop_order, self.limit_order]:
                self.martin_num = 0
                self.log(
                    'STOP/LIMIT ORDER EXECUTED, Price: %.3f, Size: %.3f, Cost: %.3f, Comm %.3f' %
                    (
                        order.executed.price,
                        order.executed.size,
                        order.executed.value,
                        order.executed.comm
                    )
                )
            # self.log_order(order)
        elif order.status in [order.Canceled, order.Rejected]:
            # self.log_order(order)
            pass
        elif order.status == order.Margin:
            # self.log_order(order)
            pass

    # def stop(self):
    #     file=open("data/resout3.csv","a")
    #     file.writelines("%s,%s,%s,%.2f \n" %
    #              (self.params.min_gap, self.p.stoploss, self.p.targetprofit, self.broker.getvalue()))
    #     file.close()
    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log('OPERATION PROFIT, MARTIN_NUM %s, GROSS %.2f, NET %.2f, ' %
                 (len(self.open_orders), trade.pnl, trade.pnlcomm))

        self.open_orders = []

    def next(self):
        if self.broker.get_cash() < 0:
            self.log("爆仓了 %s" % self.broker.get_cash())
            exit()
        self.deal_signal(self.data.volume[0])

    def deal_signal(self, signal):
        close = self.data.close[0]
        df = 0.8
        df2 = 0.0


        if signal < 1 and len(self) - self.bar_executed > 0:
            if not self.position or (
                    self.position.size > 0 and self.open_orders[-1].executed.price - close > self.params.min_gap*self.point):
                self.open_buy()

        if signal > 5 and len(self) - self.bar_executed > 0:
            if not self.position or (
                    self.position.size < 0 and close - self.open_orders[-1].executed.price > self.params.min_gap*self.point ):
                self.open_sell()

        # if 0 < self.martin_num < self.martin_num_limit:
        if signal < 2 and self.position.size < 0:
            self.close_all()

        if signal > 4 and self.position.size > 0:
            self.close_all()

    def max_buy(self):
        return self.broker.get_cash() / self.data.close[0]

    def get_size(self):

        comminfo = self.broker.getcommissioninfo(self.data)
        price = self.data.close[0]
        value = 1 * self.base_martin
        size = comminfo.getsize(price, value)
        return value

    def open_buy(self):
        close = self.data.close[0]
        size = self.get_size()
        order = self.buy(size=self.get_size(), transmit=False)
        self.open_orders.append(order)

        if self.stop_order:
            self.cancel(self.stop_order)
        if self.limit_order:
            self.cancel(self.limit_order)

        limitprice = close + self.params.stoploss * self.point
        size = self.position.size + size
        self.limit_order = self.sell(exectype=bt.Order.Limit, size=size, price=limitprice, parent=order, transmit=False)

        self.stop_order = self.sell(exectype=bt.Order.StopTrail, size=size, price=close, trailamount=self.params.targetprofit * self.point,
                                    parent=order)

        return self.open_orders

    def open_sell(self):
        close = self.data.close[0]
        size = self.get_size()
        order = self.sell(size=size, transmit=False)
        self.open_orders.append(order)

        if self.stop_order:
            self.cancel(self.stop_order)
        if self.limit_order:
            self.cancel(self.limit_order)

        limitprice = close - self.params.stoploss * self.point
        size = -self.position.size + size
        self.limit_order = self.buy(exectype=bt.Order.Limit, size=size, price=limitprice, parent=order, transmit=False)

        self.stop_order = self.buy(exectype=bt.Order.StopTrail, size=size, price=close, trailamount=self.params.targetprofit * self.point,
                                   parent=order)

        return self.open_orders

    def close_all(self):
        self.close_order = self.close()
        self.cancel(self.stop_order)
        self.cancel(self.limit_order)
